[ExecuTorch][WebGPU] Add clone op (aten.clone.default)#20463
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20463
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 54 PendingAs of commit 75023ac with merge base 200f64a ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * __->__ #20362 * #20361 * #20360 * #20359 Adds `aten.select_copy.int` to the WebGPU delegate as a gather: picks a fixed index along one dim, producing an output of rank (input rank - 1). Composition (single dispatch): - `select/Select.cpp` — reads `[self, dim, index, out]` (static `Int` via `read_scalar`; throws on dynamic `SymInt`), normalizes + bounds-checks dim/index, builds 2 `TensorMeta` UBOs + a `SelectParams{dim,index}`, fp32 guard, 1D-dispatch over `numel`, releases uniforms after the bind group. - `select/select.wgsl` — seeds the input offset with `index * in.strides[dim]`, delinearizes the output index, maps each out dim to its in dim (shifted past the selected dim), relinearizes on input strides. @exported-using-ghexport Differential Revision: [D108793166](https://our.internmc.facebook.com/intern/diff/D108793166/) Differential Revision: [D108793166](https://our.internmc.facebook.com/intern/diff/D108793166)
…ework) (#20363) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * #20393 * #20392 * #20391 * #20390 * __->__ #20363 * #20362 * #20361 * #20360 * #20359 Registers `aten.select_copy.int` in the `cases.py` op-test framework: a `_select_suite` of 4 configs (leading/middle/last dim + negative index) that `generate_op_tests` exports and compares to a torch golden on Dawn. Also adds `test/ops/select/test_select.py` (`SelectModule` + `CONFIGS` + an export-delegation/eager smoke test) and the `aten.select_copy.int` partitioner-allowlist entry in `tester.py`. @exported-using-ghexport Differential Revision: [D108793161](https://our.internmc.facebook.com/intern/diff/D108793161/) Differential Revision: [D108793161](https://our.internmc.facebook.com/intern/diff/D108793161)
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * #20393 * #20392 * #20391 * __->__ #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Adds `aten.sigmoid.default` to the WebGPU delegate: element-wise `1/(1+exp(-x))` over a flat fp32 buffer. On the Llama critical path (`F.silu` -> `sigmoid` + `mul`). Composition (single dispatch): - `sigmoid/UnaryOp.cpp` — binds input (storage, read-only) + output (storage) + a `Params{num_elements}` uniform, 1D-dispatches over `num_elements` with `override wg_size` (clamped to the device limit); mirrors the `add` op (uniform mapped-at-creation, released after the bind group). - `sigmoid/sigmoid.wgsl` — guards `idx >= num_elements` and writes the logistic of each element. @exported-using-ghexport Differential Revision: [D108793157](https://our.internmc.facebook.com/intern/diff/D108793157/) Differential Revision: [D108793157](https://our.internmc.facebook.com/intern/diff/D108793157)
…k) (#20391) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * #20393 * #20392 * __->__ #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Registers `aten.sigmoid.default` in the `cases.py` op-test framework: a `_sigmoid_suite` (hard-coded shapes + a saturation case over a `linspace(-12, 12)` input) that `generate_op_tests` exports and compares to an fp64 torch golden on Dawn. Also adds `test/ops/sigmoid/test_sigmoid.py` (`SigmoidModule` + `N` + `_det_input` + an export-delegation/eager smoke test) and the `aten.sigmoid.default` partitioner-allowlist entry in `tester.py`. @exported-using-ghexport Differential Revision: [D108793159](https://our.internmc.facebook.com/intern/diff/D108793159/) Differential Revision: [D108793159](https://our.internmc.facebook.com/intern/diff/D108793159)
…20392) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * #20393 * __->__ #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Adds `aten.squeeze_copy.dims` and `aten.unsqueeze_copy.default` to the WebGPU delegate. Both are numel-preserving shape ops; on a dense row-major buffer backend they are the same flat copy as `view_copy` — only the shape metadata differs (mirrors the Vulkan delegate, which routes both through `add_view_copy_node`). Composition (no new kernel): - `squeeze/Squeeze.cpp` — reads `args = [self, dims, out]`, ignores the AOT-fixed `dims`, calls `add_flat_copy(graph, in, out)` from `runtime/ops/view_copy/view_copy.h`. - `unsqueeze/Unsqueeze.cpp` — reads `args = [self, dim, out]`, ignores the AOT-fixed `dim`, calls `add_flat_copy(graph, in, out)`. @exported-using-ghexport Differential Revision: [D108793153](https://our.internmc.facebook.com/intern/diff/D108793153/) Differential Revision: [D108793153](https://our.internmc.facebook.com/intern/diff/D108793153)
….py op-test framework) (#20393) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * #20394 * __->__ #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Registers `aten.squeeze_copy.dims` and `aten.unsqueeze_copy.default` in the `cases.py` op-test framework: a `_squeeze_suite` of 3 configs (squeeze leading/middle/multiple size-1 dims) and a `_unsqueeze_suite` of 3 configs (insert dim at front/middle/last) that `generate_op_tests` exports via `VulkanPartitioner` and compares to a torch golden on Dawn. Also adds `test/ops/squeeze/test_squeeze.py` (`SqueezeModule` + `CONFIGS` + `_op_delegated` smoke test), `test/ops/unsqueeze/test_unsqueeze.py` (`UnsqueezeModule` + `CONFIGS` + `_op_delegated` smoke test), and the two partitioner-allowlist entries in `tester.py`. @exported-using-ghexport Differential Revision: [D108793152](https://our.internmc.facebook.com/intern/diff/D108793152/) Differential Revision: [D108793152](https://our.internmc.facebook.com/intern/diff/D108793152)
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * #20395 * __->__ #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Adds `aten.slice_copy.Tensor` to the WebGPU delegate as a gather: each output element is mapped back to its source input element along the sliced dim via `start + coord * step`. Composition (single compute dispatch): - `runtime/ops/slice/Slice.cpp` — reads `args = [self, dim, start, end, step, out]` via `read_scalar` (static `Int`/`Null`-sentinel default; throws on dynamic `SymInt`); normalizes negative `dim`/`start`, clamps `start` to `[0, in_size]`; builds two `TensorMeta` UBOs + a `SliceParams{dim, start, step}` uniform; guards fp32; dispatches over `compute_1d_workgroup_count(out.numel)` with `override wg_size`; releases all uniforms after the bind group. - `runtime/ops/slice/slice.wgsl` — delinearizes the output index over the contiguous output strides, maps the sliced-dim coordinate back to the input (`start + coord*step`), relinearizes over the input strides. @exported-using-ghexport Differential Revision: [D108793168](https://our.internmc.facebook.com/intern/diff/D108793168/) Differential Revision: [D108793168](https://our.internmc.facebook.com/intern/diff/D108793168)
…work) (#20395) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * #20396 * __->__ #20395 * #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Registers `aten.slice_copy.Tensor` in the `cases.py` op-test framework: a `_slice_suite` of 4 configs (leading-dim slice `[:,1:5]`, last-dim slice `[...,1:3]`, step-2 `[:,0:8:2]`, negative-end `[:,1:-1]`) that `generate_op_tests` exports via `VulkanPartitioner` and compares to a torch golden on Dawn. Also adds `test/ops/slice/test_slice.py` (`SliceModule` + `CONFIGS` + export-delegation/eager smoke test) and the `aten.slice_copy.Tensor` partitioner-allowlist entry in `tester.py`. @exported-using-ghexport Differential Revision: [D108793151](https://our.internmc.facebook.com/intern/diff/D108793151/) Differential Revision: [D108793151](https://our.internmc.facebook.com/intern/diff/D108793151)
…ermute_copy.default) (#20396) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * #20397 * __->__ #20396 * #20395 * #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Adds `aten.permute_copy.default` (a coordinate-reorder gather) to the WebGPU delegate, and the `IntList` graph value type it needs to read its `dims` argument. Composition: - `runtime/WebGPUGraph.{h,cpp}` — adds `ValueType::IntList` backed by `std::vector<std::vector<int64_t>> int_lists_` + `get_int_list(int)`; `build()` deserializes `vkgraph::GraphTypes::IntList` via `value_as_IntList()->items()` (int64, matching the FlatBuffer `[long]`); mirrors the existing scalar value plumbing. - `runtime/ops/permute/Permute.cpp` — reads the permutation via `get_int_list`, normalizes negative dims, validates it is a permutation of `[0, ndim)`, builds two `TensorMeta` UBOs + a `PermuteParams{perm: vec4<u32>}` uniform, guards fp32 + rank≤4, dispatches over `compute_1d_workgroup_count(out.numel)` with `override wg_size`; releases all uniforms after the bind group. - `runtime/ops/permute/permute.wgsl` — delinearizes the output index over the contiguous output strides, reads `input` at `in.strides[perm[d]]` per dim (mirrors Vulkan `permute_buffer.glsl`). - Registers both `aten.permute_copy.default` and `aten.permute.default` to the same handler. @exported-using-ghexport Differential Revision: [D108793162](https://our.internmc.facebook.com/intern/diff/D108793162/) Differential Revision: [D108793162](https://our.internmc.facebook.com/intern/diff/D108793162)
…mework) (#20397) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * #20398 * __->__ #20397 * #20396 * #20395 * #20394 * #20393 * #20392 * #20391 * #20390 * #20363 * #20362 * #20361 * #20360 * #20359 Registers `aten.permute_copy.default` in the `cases.py` op-test framework: a `_permute_suite` of 4 configs (3D rotation, 4D middle-dim transpose, 2D transpose, full 4D shuffle) that `generate_op_tests` exports via `VulkanPartitioner` and compares to a torch golden on Dawn. Also adds `test/ops/permute/test_permute.py` (`PermuteModule` + `CONFIGS` + `_op_delegated` smoke test) and the `aten.permute_copy.default` partitioner-allowlist entry in `tester.py`. @exported-using-ghexport Differential Revision: [D108793156](https://our.internmc.facebook.com/intern/diff/D108793156/) Differential Revision: [D108793156](https://our.internmc.facebook.com/intern/diff/D108793156)
…efault) (#20398) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * #20399 * __->__ #20398 Adds `aten.cat.default` to the WebGPU delegate as an index-math scatter, and the `ValueList` graph value type it needs to read its tensor-list argument. Composition (one dispatch per input): - `runtime/WebGPUGraph.{h,cpp}` — adds `ValueType::ValueList` backed by `std::vector<std::vector<int>> value_lists_` + `get_value_list(int)`; `build()` deserializes `vkgraph::GraphTypes::ValueList` via `value_as_ValueList()->items()`; mirrors the existing scalar value plumbing. - `runtime/ops/cat/Cat.cpp` — reads the input-tensor list via `get_value_list`, reads the static `Int` dim, validates each input rank + non-concat dims against the output; builds one shared `out_meta` `TensorMeta` uniform + a shared bind-group/pipeline layout; per input builds a fresh pipeline + bind group with `in_meta` + `CatParams{concat_dim, off_k}`, dispatches over `compute_1d_workgroup_count(in.numel)`; releases per-input uniforms after each bind group, the shared `out_meta` after the loop. - `runtime/ops/cat/cat.wgsl` — delinearizes the input index over the input strides, shifts the concat-dim coordinate by the host-computed `off_k` (running sum of prior input sizes), relinearizes over the output strides to scatter `output[out] = input[in]`. - The N disjoint-slab writes share one output buffer across separate dispatches and rely on the existing per-pass `execute()` ordering. @exported-using-ghexport Differential Revision: [D108793165](https://our.internmc.facebook.com/intern/diff/D108793165/) Differential Revision: [D108793165](https://our.internmc.facebook.com/intern/diff/D108793165)
…20399) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * #20435 * __->__ #20399 * #20398 Registers `aten.cat.default` in the `cases.py` op-test framework: a `_cat_suite` of 4 configs (3 inputs along dim 0, 2 inputs along dim 1, 3 inputs along dim 2, uneven split along dim 1) that `generate_op_tests` exports via `VulkanPartitioner` and compares to a torch golden on Dawn. Also adds `test/ops/cat/test_cat.py` (`CatModule` + `CONFIGS` + `_op_delegated` smoke test, with distinct per-input value ranges to catch cross-slab contamination) and the `aten.cat.default` partitioner-allowlist entry in `tester.py`. @exported-using-ghexport Differential Revision: [D108793163](https://our.internmc.facebook.com/intern/diff/D108793163/) Differential Revision: [D108793163](https://our.internmc.facebook.com/intern/diff/D108793163)
…>.py (#20435) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * #20464 * #20463 * __->__ #20435 * #20399 * #20398 Pure test-file relocation: moves the already-landed ops' tests from nested `test/ops/<op>/test_<op>.py` to flat `test/ops/test_<op>.py`, matching the ExecuTorch convention (XNNPACK uses flat `test/ops/test_<op>.py`; Vulkan uses flat `test/test_*.py`) and completing the flatten applied to the new ops in the stack below. Drops the per-op `__init__.py`; the parent `test/ops/__init__.py` is kept. Ops: `add`, `rms_norm`, `sdpa` (`test_sdpa` + `test_update_cache`), `dispatch_order`, `quantized_linear`, `embedding_q4gsw`, `rope`, `prepack`. No behavior change — the test modules and their export/golden functions are unchanged; only their path moves. Every reference to the old paths is updated: the `cases.py` op-test imports (`add`, `rms_norm`), `test/TARGETS` (`test_add` srcs), `test/ops/test_dispatch_order.py`'s internal `rms_norm` import, and the build/CI scripts that import the per-op export functions (`test/test_build_webgpu.sh`, `scripts/test_webgpu_native_ci.sh`). Nothing required the per-op subdirectory: the codegen framework imports only `cases.py`, the one buck target uses a literal path, and the native-golden scripts import the modules by path — each resolves identically at the flat path. @exported-using-ghexport Differential Revision: [D109349894](https://our.internmc.facebook.com/intern/diff/D109349894/) Differential Revision: [D109349894](https://our.internmc.facebook.com/intern/diff/D109349894)
c6b5557
into
gh/JulianCloudNTH/58/base
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #20465 * __->__ #20464 * #20463 * #20435 * #20399 * #20398 Adds the WebGPU delegate handler for aten.index.Tensor, the 1D-self advanced-index gather out[i] = self[index[i]] (output shape == index shape). This is the form the VulkanPartitioner delegates -- it requires a 1D self and exactly one non-None index (op_registry.py); 2D mask/freqs gathers stay on CPU. It mirrors the Vulkan delegate's index_tensor op (IndexTensor.cpp + index_tensor_buffer.glsl) as a single compute dispatch over the output elements, each reading the int32 index and gathering the corresponding fp32 self element. The op is composed as: - index.wgsl: one workgroup-strided pass, out[i] = self[u32(index[i])], guarded by a numel bound; buffer-only, fp32 self/out, int32 index, 1D dispatch via the shared WebGPUUtils helpers (clamp workgroup size + 1D count). - Index.cpp: validates the args (self/out tensors; indices ValueList with exactly one index tensor; fp32 self/out; int32 index; out numel == index numel), failing loud on any violation, then records the dispatch. row_width is dropped (always 1 for 1D self). @exported-using-ghexport Differential Revision: [D109478967](https://our.internmc.facebook.com/intern/diff/D109478967/) Differential Revision: [D109478967](https://our.internmc.facebook.com/intern/diff/D109478967)
…lden) (#20465) Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * __->__ #20465 * #20464 * #20463 * #20435 * #20399 * #20398 Adds the test suite for the aten.index.Tensor op (stacked on the op diff): - test/ops/index/test_index.py: exports a module computing x[idx] through VulkanPartitioner for four configs (reorder/repeat indices over distinct self values, so a wrong-gather is visible), asserts a VulkanBackend delegate with index.Tensor absorbed (not a CPU fallback), and writes per-config .pte + .self/.idx/.golden.bin. - test/native/test_index.cpp: a standalone Dawn binary that loads each .pte, feeds self (fp32) + index (int64 at the program boundary, narrowed to the int32 buffer) and compares the gather against the torch golden at 1e-3, with a single-output shape guard. - Wired into CMake (webgpu_index_test), test/TARGETS (python_unittest test_index), and the Dawn native CI script. @exported-using-ghexport Differential Revision: [D109479000](https://our.internmc.facebook.com/intern/diff/D109479000/) Differential Revision: [D109479000](https://our.internmc.facebook.com/intern/diff/D109479000)
Pull Request resolved: #20463 `aten.clone.default` is a pure flat copy on the buffer-only WebGPU backend, identical to `view_copy`: `clone_impl` reuses the existing `add_flat_copy` helper (`output[i] = input[i]`) and registers a handler under `aten.clone.default`. No new shader, generated WGSL header, or CMake source — it shares the `view_copy` flat-copy compute pipeline. Required for end-to-end Llama 3.2 1B (4-bit, KV cache): the exported model serializes 2 `aten.clone.default` ops into its runtime operator chain (the RoPE-frequency clones reused across all 16 transformer layers), so without a handler the partition graph-breaks at those nodes. Mirrors the Vulkan delegate, which registers the same op and routes a buffer clone to a flat view-copy. ghstack-source-id: 397534700 @exported-using-ghexport @diff-train-skip-merge Differential Revision: [D109477717](https://our.internmc.facebook.com/intern/diff/D109477717/)
Pull Request resolved: #20463 `aten.clone.default` is a pure flat copy on the buffer-only WebGPU backend, identical to `view_copy`: `clone_impl` reuses the existing `add_flat_copy` helper (`output[i] = input[i]`) and registers a handler under `aten.clone.default`. No new shader, generated WGSL header, or CMake source — it shares the `view_copy` flat-copy compute pipeline. Required for end-to-end Llama 3.2 1B (4-bit, KV cache): the exported model serializes 2 `aten.clone.default` ops into its runtime operator chain (the RoPE-frequency clones reused across all 16 transformer layers), so without a handler the partition graph-breaks at those nodes. Mirrors the Vulkan delegate, which registers the same op and routes a buffer clone to a flat view-copy. ghstack-source-id: 397534700 @exported-using-ghexport @diff-train-skip-merge Differential Revision: [D109477717](https://our.internmc.facebook.com/intern/diff/D109477717/)
Pull Request resolved: #20463 `aten.clone.default` is a pure flat copy on the buffer-only WebGPU backend, identical to `view_copy`: `clone_impl` reuses the existing `add_flat_copy` helper (`output[i] = input[i]`) and registers a handler under `aten.clone.default`. No new shader, generated WGSL header, or CMake source — it shares the `view_copy` flat-copy compute pipeline. Required for end-to-end Llama 3.2 1B (4-bit, KV cache): the exported model serializes 2 `aten.clone.default` ops into its runtime operator chain (the RoPE-frequency clones reused across all 16 transformer layers), so without a handler the partition graph-breaks at those nodes. Mirrors the Vulkan delegate, which registers the same op and routes a buffer clone to a flat view-copy. ghstack-source-id: 397534700 @exported-using-ghexport @diff-train-skip-merge Differential Revision: [D109477717](https://our.internmc.facebook.com/intern/diff/D109477717/)
Stack from ghstack (oldest at bottom):
aten.clone.defaultis a pure flat copy on the buffer-only WebGPU backend, identical toview_copy:clone_implreuses the existingadd_flat_copyhelper (output[i] = input[i]) and registers a handler underaten.clone.default. No new shader, generated WGSL header, or CMake source — it shares theview_copyflat-copy compute pipeline. Required for end-to-end Llama 3.2 1B (4-bit, KV cache): the exported model serializes 2aten.clone.defaultops into its runtime operator chain (the RoPE-frequency clones reused across all 16 transformer layers), so without a handler the partition graph-breaks at those nodes. Mirrors the Vulkan delegate, which registers the same op and routes a buffer clone to a flat view-copy.@exported-using-ghexport
Differential Revision: D109477717
Differential Revision: D109477717